// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <limits>
#include <tbir/runtime/container.h>

namespace tbir::runtime {

    static string_view PREFIX_KEY("____prefix_path____");

    static constexpr int NONE_DEVICE = INT16_MIN;

    struct Attributes {
    public:
        bool HasAttr(string_view key) {
            auto itr = attrs_.find(key);
            return itr != attrs_.end();
        }

        template<class U>
        U GetAttr(string_view key, const U &default_val = U{}) const {
            using U_TYPE = typename std::remove_cv<typename std::remove_reference<U>::type>::type;
            GenericValueConverter<U_TYPE> Converter;
            auto itr = attrs_.find(key);
            if (itr != attrs_.end()) {
                return Converter(itr->second);
            } else {
                return default_val;
            }
        }

        template<class U>
        void SetAttr(string_view key, U &&val) {
            using Converter = GenericValueConverter<RTValue>;
            attrs_[String(key.data(), key.size())] = Converter()(std::forward<U>(val));
        }

        static Attributes FromDict(const Dict &generic_attrs) {
            Attributes res_attrs;
            for (auto &attr : generic_attrs.items()) {
                MXCHECK(attr.first.type_code() == TypeIndex::kRuntimeString ||
                        attr.first.type_code() == TypeIndex::kRuntimeUnicode);
                if (attr.first.type_code() == TypeIndex::kRuntimeString) {
                    res_attrs.attrs_[attr.first.As<String>()] = attr.second;
                } else {
                    res_attrs.attrs_[attr.first.As<Unicode>().encode()] = attr.second;
                }
            }
            return res_attrs;
        }

        Dict ToDict() const {
            Dict d;
            d.reserve(attrs_.size());
            for (auto &attr : attrs_) {
                d[attr.first] = attr.second;
            }
            return d;
        }

    protected:
        bool profiling = false;
        ska::flat_hash_map<String, RTValue> attrs_;

        friend class OpKernel;

        friend class TXSession;
    };

    template<>
    inline String Attributes::GetAttr(string_view key, const String &default_val) const {
        auto itr = attrs_.find(key);
        if (itr != attrs_.end()) {
            if (itr->second.type_code() == TypeIndex::kRuntimeUnicode) {
                return UnicodeHelper::Encode(itr->second.AsNoCheck<unicode_view>());
            }
            return itr->second.As<String>();
        } else {
            return default_val;
        }
    }

    template<>
    inline Unicode Attributes::GetAttr(string_view key, const Unicode &default_val) const {
        auto itr = attrs_.find(key);
        if (itr != attrs_.end()) {
            if (itr->second.type_code() == TypeIndex::kRuntimeString) {
                return StringHelper::Decode(itr->second.AsNoCheck<string_view>());
            }
            return itr->second.As<Unicode>();
        } else {
            return default_val;
        }
    }

}  // namespace tbir::runtime
